#include "GRU.h"

#include "activation.h"

GRUCell::GRUCell(const int hidden_size)
    : _hidden_size(hidden_size)
    , _hidden(hidden_size, 1)
    , _w_ir(hidden_size, 1)
    , _w_iz(hidden_size, 1)
    , _w_in(hidden_size, 1)
    , _w_hr(hidden_size, hidden_size)
    , _w_hz(hidden_size, hidden_size)
    , _w_hn(hidden_size, hidden_size)
    , _b_ir(hidden_size)
    , _b_iz(hidden_size)
    , _b_in(hidden_size)
    , _b_hr(hidden_size)
    , _b_hz(hidden_size)
    , _b_hn(hidden_size)
    , _r(hidden_size)
    , _z(hidden_size)
    , _n(hidden_size)
{
}

void GRUCell::process(const Eigen::VectorXf& input, Eigen::VectorXf& output)
{
    Eigen::VectorXf rt(_hidden_size), zt(_hidden_size), nt(_hidden_size),
        hadamard(_hidden_size);
    /*  Math:
      R = sigmoid(Wir*X + Bir + Whr*Hidden1 + Bhr)
      Z = sigmoid(Wiz*X + Biz + Whz*Hidden1 + Bhz)
      N = tanh(Win*X + Bin + hadmard(R, (Whn*Hidden1 + Bhn)))
      Hidden = hadmard((1-Z), N) + hadmard(Z, Hidden1)
  */
    rt = _w_ir * input + _b_ir + _w_hr * _hidden + _b_hr;
    activation::sigmoid(rt, _r);
    zt = _w_iz * input + _b_iz + _w_hz * _hidden + _b_hz;
    activation::sigmoid(zt, _z);
    nt = _w_in * input + _b_in;
    hadamard = _r.array() * (_w_hn * _hidden + _b_hn).array();
    nt += hadamard;
    activation::tanh(nt, _n);
    _hidden = (Eigen::VectorXf::Ones(_hidden_size) - _z).array() * _n.array() + _z.array() * _hidden.array();
    output = _hidden;
}

void GRUCell::reset()
{
    for (int i = 0; i < _hidden_size; i++) {
        _hidden(i, 0) = 0.f;
        _r(i) = 0.f;
        _z(i) = 0.f;
        _n(i) = 0.f;
    }
}

void GRUCell::setParams(float* w_ir, float* w_iz, float* w_in, float* w_hr,
    float* w_hz, float* w_hn, float* b_ir, float* b_iz,
    float* b_in, float* b_hr, float* b_hz, float* b_hn)
{
    for (int i = 0; i < _hidden_size; i++) {
        _w_ir(i, 0) = w_ir[i];
        _w_iz(i, 0) = w_iz[i];
        _w_in(i, 0) = w_in[i];
        for (int j = 0; j < _hidden_size; j++) {
            _w_hr(i, j) = w_hr[i * _hidden_size + j];
            _w_hz(i, j) = w_hz[i * _hidden_size + j];
            _w_hn(i, j) = w_hn[i * _hidden_size + j];
        }
        _b_ir(i) = b_ir[i];
        _b_iz(i) = b_iz[i];
        _b_in(i) = b_in[i];
        _b_hr(i) = b_hr[i];
        _b_hz(i) = b_hz[i];
        _b_hn(i) = b_hn[i];
    }
}

GRU::GRU(const int hidden_size)
    : _hidden_size(hidden_size)
    , _head_weight(hidden_size)
    , _head_bias(1)
    , _input(1)
    , _hidden(hidden_size)
    , _output(1)
{
    _rec = GRUCell(hidden_size);
    _lin = Linear(hidden_size);
}

void GRU::process(const float* input, float* output, size_t num_samples)
{
    for (int i = 0; i < num_samples; i++) {
        output[i] = _processSample(input[i]);
    }
}

float GRU::_processSample(const float x)
{
    _input(0) = x;
    _rec.process(_input, _hidden);
    _lin.process(_hidden, _output);

    return _output(0);
}

void GRU::setParams(const std::string name, void* params)
{
    if (name == "Linear") {
        LinearInitStruct_t* linParams = (LinearInitStruct_t*)params;
        _lin.setParams(linParams->w, linParams->b);
    } else if (name == "GRU") {
        GruInitStruct_t* gruParams = (GruInitStruct_t*)params;
        _rec.setParams(gruParams->w_ir, gruParams->w_iz, gruParams->w_in,
            gruParams->w_hr, gruParams->w_hz, gruParams->w_hn,
            gruParams->b_ir, gruParams->b_iz, gruParams->b_in,
            gruParams->b_hr, gruParams->b_hz, gruParams->b_hn);
    }
}

void GRU::reset() 
{ 
    _rec.reset();
    _input(0) = 0.f;
    for (int i = 0; i < _hidden_size; i++) {
        _hidden(i) = 0.f;
    }
    _output(0) = 0.f;
}
